package com.dtdream.emr;

import org.apache.hadoop.hive.ql.exec.UDFArgumentException;
import org.apache.hadoop.hive.ql.metadata.HiveException;
import org.apache.hadoop.hive.ql.udf.generic.GenericUDTF;
import org.apache.hadoop.hive.serde2.objectinspector.ObjectInspector;
import org.apache.hadoop.hive.serde2.objectinspector.ObjectInspectorFactory;
import org.apache.hadoop.hive.serde2.objectinspector.StructObjectInspector;
import org.apache.hadoop.hive.serde2.objectinspector.primitive.DoubleObjectInspector;
import org.apache.hadoop.hive.serde2.objectinspector.primitive.LongObjectInspector;
import org.apache.hadoop.hive.serde2.objectinspector.primitive.PrimitiveObjectInspectorFactory;

import java.util.ArrayList;
import java.util.List;

/**
 * @author xukai
 * @desc
 * @date 2018-09-07-9:27
 */
public class VectorDistance extends GenericUDTF {

    private LongObjectInspector longOI = null;

    private DoubleObjectInspector doubleOI = null;

    @Override
    public StructObjectInspector initialize(ObjectInspector[] args) throws UDFArgumentException {
        if (args.length != 102) {
            throw new UDFArgumentException("VectorDistance() takes exactly 102 argument");
        }


        // 输入格式（inspectors）
        longOI = (LongObjectInspector) args[0];
        doubleOI = (DoubleObjectInspector) args[1];

        // 输出格式（inspectors） -- 有两个属性的对象
        List<String> fieldNames = new ArrayList<String>(2);
        List<ObjectInspector> fieldOIs = new ArrayList<ObjectInspector>(3);
        fieldNames.add("id");
        fieldNames.add("distance");
        fieldOIs.add(PrimitiveObjectInspectorFactory.javaLongObjectInspector);
        fieldOIs.add(PrimitiveObjectInspectorFactory.javaDoubleObjectInspector);
        return ObjectInspectorFactory.getStandardStructObjectInspector(fieldNames, fieldOIs);
    }

    @Override
    public void process(Object[] objects) throws HiveException {
        Long id = (Long)(longOI.getPrimitiveJavaObject(objects[0]));
        if (id == null || "".equals(id)) {
            return;
        }

        double sum = 0;
        for (int i = 0; i < 50; i++) {
            double v1 = (double)doubleOI.getPrimitiveJavaObject(objects[i + 1]);
            double v2 = (double)doubleOI.getPrimitiveJavaObject(objects[i + 52]);
            sum += Math.pow((v1 - v2), 2);
        }

        forward(new Object[]{id, Math.sqrt(sum)});
    }

    @Override
    public void close() throws HiveException {

    }
}
